{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%env LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1\n",
    "%env LLM_API_KEY=sk-替换为自己的Qwen API Key"
   ],
   "id": "cce596b168e32c4c"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "af375836-b870-458b-87d1-4e00565977eb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:05.075435Z",
     "iopub.status.busy": "2024-11-24T07:22:05.075303Z",
     "iopub.status.idle": "2024-11-24T07:22:05.087798Z",
     "shell.execute_reply": "2024-11-24T07:22:05.087352Z",
     "shell.execute_reply.started": "2024-11-24T07:22:05.075422Z"
    },
    "papermill": {
     "duration": 0.115454,
     "end_time": "2024-11-23T14:29:00.919641",
     "exception": false,
     "start_time": "2024-11-23T14:29:00.804187",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "!pip install -U langchain langchain-community langchain-openai langchain-cohere pypdf sentence_transformers chromadb shutil openpyxl FlagEmbedding cohere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1e2c72b8-ee12-4130-af88-699998aa230c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:05.088372Z",
     "iopub.status.busy": "2024-11-24T07:22:05.088248Z",
     "iopub.status.idle": "2024-11-24T07:22:05.300592Z",
     "shell.execute_reply": "2024-11-24T07:22:05.300021Z",
     "shell.execute_reply.started": "2024-11-24T07:22:05.088360Z"
    },
    "papermill": {
     "duration": 0.319981,
     "end_time": "2024-11-23T14:29:01.380771",
     "exception": false,
     "start_time": "2024-11-23T14:29:01.060790",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "841d2b02-ad06-40d2-b11f-c7adccec6ca2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:05.301365Z",
     "iopub.status.busy": "2024-11-24T07:22:05.301194Z",
     "iopub.status.idle": "2024-11-24T07:22:05.304213Z",
     "shell.execute_reply": "2024-11-24T07:22:05.303862Z",
     "shell.execute_reply.started": "2024-11-24T07:22:05.301352Z"
    },
    "papermill": {
     "duration": 0.121409,
     "end_time": "2024-11-23T14:29:01.638126",
     "exception": false,
     "start_time": "2024-11-23T14:29:01.516717",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "expr_version = 'retrieval_v13_contextual_embeddings'\n",
    "\n",
    "preprocess_output_dir = os.path.join(os.path.pardir, 'outputs', 'v1_20240713')\n",
    "expr_dir = os.path.join(os.path.pardir, 'experiments', expr_version)\n",
    "\n",
    "os.makedirs(expr_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf7e81e3-4c82-4842-aef5-7592caaf1d39",
   "metadata": {
    "papermill": {
     "duration": 0.100379,
     "end_time": "2024-11-23T14:29:01.862379",
     "exception": false,
     "start_time": "2024-11-23T14:29:01.762000",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# 读取文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e6920e29-bc7d-4635-be06-d151eaf0e100",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:05.304792Z",
     "iopub.status.busy": "2024-11-24T07:22:05.304663Z",
     "iopub.status.idle": "2024-11-24T07:22:07.061126Z",
     "shell.execute_reply": "2024-11-24T07:22:07.060650Z",
     "shell.execute_reply.started": "2024-11-24T07:22:05.304779Z"
    },
    "papermill": {
     "duration": 2.012298,
     "end_time": "2024-11-23T14:29:03.974974",
     "exception": false,
     "start_time": "2024-11-23T14:29:01.962676",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "\n",
    "loader = PyPDFLoader(os.path.join(os.path.pardir, 'data', '2024全球经济金融展望报告.pdf'))\n",
    "documents = loader.load()\n",
    "\n",
    "qa_df = pd.read_excel(os.path.join(preprocess_output_dir, 'question_answer.xlsx'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "841ec659-4ad7-4e1f-b1ea-3477bf97fde3",
   "metadata": {
    "papermill": {
     "duration": 0.100297,
     "end_time": "2024-11-23T14:29:04.219302",
     "exception": false,
     "start_time": "2024-11-23T14:29:04.119005",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# 文档切分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "74fe856a-7c19-4c3c-bb30-7abfa6298f74",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:07.061949Z",
     "iopub.status.busy": "2024-11-24T07:22:07.061623Z",
     "iopub.status.idle": "2024-11-24T07:22:07.069092Z",
     "shell.execute_reply": "2024-11-24T07:22:07.068654Z",
     "shell.execute_reply.started": "2024-11-24T07:22:07.061936Z"
    },
    "papermill": {
     "duration": 0.109229,
     "end_time": "2024-11-23T14:29:04.429069",
     "exception": false,
     "start_time": "2024-11-23T14:29:04.319840",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "def split_docs(documents, filepath, chunk_size=400, chunk_overlap=40, seperators=['\\n\\n\\n', '\\n\\n'], force_split=False):\n",
    "    if os.path.exists(filepath) and not force_split:\n",
    "        print('found cache, restoring...')\n",
    "        return pickle.load(open(filepath, 'rb'))\n",
    "\n",
    "    splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=chunk_size,\n",
    "        chunk_overlap=chunk_overlap,\n",
    "        separators=seperators\n",
    "    )\n",
    "    split_docs = splitter.split_documents(documents)\n",
    "    for chunk in split_docs:\n",
    "        chunk.metadata['uuid'] = str(uuid4())\n",
    "\n",
    "    pickle.dump(split_docs, open(filepath, 'wb'))\n",
    "\n",
    "    return split_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aa25540d-0504-4ae7-9804-9e3862b132d5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:07.070488Z",
     "iopub.status.busy": "2024-11-24T07:22:07.070358Z",
     "iopub.status.idle": "2024-11-24T07:22:07.079493Z",
     "shell.execute_reply": "2024-11-24T07:22:07.079049Z",
     "shell.execute_reply.started": "2024-11-24T07:22:07.070476Z"
    },
    "papermill": {
     "duration": 0.145583,
     "end_time": "2024-11-23T14:29:04.677429",
     "exception": false,
     "start_time": "2024-11-23T14:29:04.531846",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found cache, restoring...\n"
     ]
    }
   ],
   "source": [
    "splitted_docs = split_docs(documents, os.path.join(preprocess_output_dir, 'split_docs.pkl'), chunk_size=500, chunk_overlap=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "21ae2e70-72f3-44e4-b1af-85348e8e70c5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:07.080057Z",
     "iopub.status.busy": "2024-11-24T07:22:07.079934Z",
     "iopub.status.idle": "2024-11-24T07:22:07.201622Z",
     "shell.execute_reply": "2024-11-24T07:22:07.201128Z",
     "shell.execute_reply.started": "2024-11-24T07:22:07.080045Z"
    },
    "papermill": {
     "duration": 0.255459,
     "end_time": "2024-11-23T14:29:05.065932",
     "exception": false,
     "start_time": "2024-11-23T14:29:04.810473",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import re\n",
    "from langchain.schema import Document\n",
    "\n",
    "# 去除页眉\n",
    "pattern = r\"^全球经济金融展望报告\\n中国银行研究院 \\d+ 2024年\"\n",
    "doc_content = '\\n'.join(re.sub(pattern, '', doc.page_content) for doc in documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "220dbc3a-fceb-4e49-a3f1-01e16660b2a6",
   "metadata": {
    "papermill": {
     "duration": 0.100209,
     "end_time": "2024-11-23T14:29:05.255871",
     "exception": false,
     "start_time": "2024-11-23T14:29:05.155662",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# 检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8598a11c-25d8-4af1-a98b-06a8c394e261",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:07.202335Z",
     "iopub.status.busy": "2024-11-24T07:22:07.202113Z",
     "iopub.status.idle": "2024-11-24T07:22:08.036594Z",
     "shell.execute_reply": "2024-11-24T07:22:08.036104Z",
     "shell.execute_reply.started": "2024-11-24T07:22:07.202322Z"
    },
    "papermill": {
     "duration": 0.989203,
     "end_time": "2024-11-23T14:29:06.345534",
     "exception": false,
     "start_time": "2024-11-23T14:29:05.356331",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cuda\n"
     ]
    }
   ],
   "source": [
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "import torch\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# device = 'cpu'\n",
    "print(f'device: {device}')\n",
    "\n",
    "def get_embeddings(model_path):\n",
    "    embeddings = HuggingFaceBgeEmbeddings(\n",
    "        model_name=model_path,\n",
    "        model_kwargs={'device': device},\n",
    "        encode_kwargs={'normalize_embeddings': True},\n",
    "        # show_progress=True\n",
    "        query_instruction='为这个句子生成表示以用于检索相关文章：'\n",
    "    )\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e45931bb-6812-4680-9e73-2c2e63a7bbcb",
   "metadata": {
    "papermill": {
     "duration": 0.111733,
     "end_time": "2024-11-23T14:29:06.613089",
     "exception": false,
     "start_time": "2024-11-23T14:29:06.501356",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 构建Contextual Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "52174650-948f-4b11-9292-b729e65a97ea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:08.037305Z",
     "iopub.status.busy": "2024-11-24T07:22:08.037075Z",
     "iopub.status.idle": "2024-11-24T07:22:08.185432Z",
     "shell.execute_reply": "2024-11-24T07:22:08.184954Z",
     "shell.execute_reply.started": "2024-11-24T07:22:08.037290Z"
    },
    "papermill": {
     "duration": 0.402828,
     "end_time": "2024-11-23T14:29:07.116501",
     "exception": false,
     "start_time": "2024-11-23T14:29:06.713673",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOllama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d20b868a-2fdc-4f11-9a8f-14286a5857de",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:08.186327Z",
     "iopub.status.busy": "2024-11-24T07:22:08.185973Z",
     "iopub.status.idle": "2024-11-24T07:22:08.195769Z",
     "shell.execute_reply": "2024-11-24T07:22:08.195323Z",
     "shell.execute_reply.started": "2024-11-24T07:22:08.186312Z"
    },
    "papermill": {
     "duration": 0.114426,
     "end_time": "2024-11-23T14:29:07.361813",
     "exception": false,
     "start_time": "2024-11-23T14:29:07.247387",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3007024/2314240529.py:1: LangChainDeprecationWarning: The class `ChatOllama` was deprecated in LangChain 0.3.1 and will be removed in 1.0.0. An updated version of the class exists in the :class:`~langchain-ollama package and should be used instead. To use it run `pip install -U :class:`~langchain-ollama` and import as `from :class:`~langchain_ollama import ChatOllama``.\n",
      "  llm = ChatOllama(\n"
     ]
    }
   ],
   "source": [
    "llm = ChatOllama(\n",
    "    model='qwen2:7b-instruct-32k',\n",
    "    base_url='http://localhost:11434',\n",
    "    top_k=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "26eaae03-f566-4358-b0ab-d2812993d743",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:08.196364Z",
     "iopub.status.busy": "2024-11-24T07:22:08.196236Z",
     "iopub.status.idle": "2024-11-24T07:22:08.202513Z",
     "shell.execute_reply": "2024-11-24T07:22:08.202104Z",
     "shell.execute_reply.started": "2024-11-24T07:22:08.196351Z"
    },
    "papermill": {
     "duration": 0.163967,
     "end_time": "2024-11-23T14:29:07.635025",
     "exception": false,
     "start_time": "2024-11-23T14:29:07.471058",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "DOCUMENT_CONTEXT_PROMPT = \"\"\"\n",
    "<document>\n",
    "{doc_content}\n",
    "</document>\n",
    "\"\"\"\n",
    "\n",
    "CHUNK_CONTEXT_PROMPT = \"\"\"\n",
    "这是我们想要放置在整个文档中的片段。\n",
    "<chunk>\n",
    "{chunk_content}\n",
    "</chunk>\n",
    "\n",
    "请为了提高检索效率，给出一个简短而精确的上下文来定位这个片段在整个文件中的位置。\n",
    "\"\"\"\n",
    "\n",
    "def situate_context(doc: str, chunk: str) -> str:\n",
    "    messages = [\n",
    "        (\"user\", DOCUMENT_CONTEXT_PROMPT.format(doc_content=doc)),\n",
    "        (\"user\", CHUNK_CONTEXT_PROMPT.format(chunk_content=chunk)),\n",
    "    ]\n",
    "    return llm.invoke(messages).content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b614a22a-ee62-412d-8320-69b56708f765",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:22:08.203098Z",
     "iopub.status.busy": "2024-11-24T07:22:08.202965Z",
     "iopub.status.idle": "2024-11-24T07:23:04.321305Z",
     "shell.execute_reply": "2024-11-24T07:23:04.320857Z",
     "shell.execute_reply.started": "2024-11-24T07:22:08.203086Z"
    },
    "papermill": {
     "duration": 7.095007,
     "end_time": "2024-11-23T14:29:14.833418",
     "exception": false,
     "start_time": "2024-11-23T14:29:07.738411",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'这个片段位于文档的开头部分。它概述了2023年全球经济金融的基本情况和2024年的展望，并提到了报告的主要作者、联系信息以及数据来源（IMF和中国银行研究院）。此外，该片段还简要介绍了全球主要经济体GDP增速的变化趋势。因此，这个片段可以作为文档的引言部分，用于概述研究的重点内容和背景信息。\\n\\n上下文定位：\\n- **位置**：文档开头\\n- **作用**：提供2023年全球经济金融概览、报告作者与联系方式、数据来源以及主要经济体GDP增速变化趋势的简要介绍。\\n- **重要性**：为读者提供了对报告核心内容的快速了解，帮助他们理解研究的主要发现和背景。\\n\\n通过这个片段，我们可以定位到文档的引言部分，它通常用于概述研究的目的、方法、关键发现等，并提供必要的上下文信息。'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "situate_context(doc_content, chunk=splitted_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "473c2187-ae9a-4223-b1f6-e54008b40149",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:23:04.321923Z",
     "iopub.status.busy": "2024-11-24T07:23:04.321798Z",
     "iopub.status.idle": "2024-11-24T07:38:53.300334Z",
     "shell.execute_reply": "2024-11-24T07:38:53.299993Z",
     "shell.execute_reply.started": "2024-11-24T07:23:04.321911Z"
    },
    "papermill": {
     "duration": 266.750722,
     "end_time": "2024-11-23T14:33:41.712280",
     "exception": false,
     "start_time": "2024-11-23T14:29:14.961558",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d28c623bcbf4b979cb1fb54c22c0ce7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/52 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "contextualized_chunks = []\n",
    "for chunk in tqdm(splitted_docs):\n",
    "    contextualized_text = situate_context(doc_content, chunk)\n",
    "    contextualized_chunks.append(\n",
    "        Document(\n",
    "            page_content=chunk.page_content + '\\n\\n' + contextualized_text,\n",
    "            metadata=chunk.metadata\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d5ae94fd-fb02-4f67-b489-298fb7b2f48f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:38:53.300933Z",
     "iopub.status.busy": "2024-11-24T07:38:53.300800Z",
     "iopub.status.idle": "2024-11-24T07:38:53.304260Z",
     "shell.execute_reply": "2024-11-24T07:38:53.303836Z",
     "shell.execute_reply.started": "2024-11-24T07:38:53.300920Z"
    },
    "papermill": {
     "duration": 0.121981,
     "end_time": "2024-11-23T14:33:41.973812",
     "exception": false,
     "start_time": "2024-11-23T14:33:41.851831",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(metadata={'source': 'data/2024全球经济金融展望报告.pdf', 'page': 0, 'uuid': 'e73a0c9d-d42b-4350-a4c3-b38bf67c68a5'}, page_content='研究院\\n全球经济金融展望报告\\n要点2024年年报（总第57期） 报告日期：2023年12月12日\\n●2023年全球经济增长动力持续回落，各国复苏分化，\\n发达经济体增速明显放缓，新兴经济体整体表现稳定。\\n全球贸易增长乏力，各国生产景气度逐渐回落，内需\\n对经济的拉动作用减弱。欧美央行货币政策紧缩态势\\n放缓，美元指数高位震荡后走弱，全球股市表现总体\\n好于预期，但区域分化明显。高利率环境抑制债券融\\n资需求，债券违约风险持续上升。\\n●展望2024年，预计全球经济复苏将依旧疲软，主要\\n经济体增长态势和货币政策走势将进一步分化。欧美\\n央行大概率结束本轮紧缩货币周期，美元指数将逐步\\n走弱，流向新兴经济体的跨境资本将增加。国际原油\\n市场短缺格局或延续，新能源发展成为重点。\\n●海湾六国经济发展与投资前景、高利率和高债务对\\n美国房地产市场脆弱性的影响等热点问题值得关注。中国银行研究院\\n全球经济金融研究课题组\\n组长：陈卫东\\n副组长：钟红\\n廖淑萍\\n成员：边卫红\\n熊启跃\\n王有鑫\\n曹鸿宇\\n李颖婷\\n王宁远\\n初晓\\n章凯莉\\n黄小军（纽约）\\n陆晓明（纽约）\\n黄承煜（纽约）\\n宋达志（伦敦）\\n李振龙（伦敦）\\n张传捷（伦敦）\\n刘冰彦（法兰克福）\\n温颍坤（法兰克福）\\n张明捷（法兰克福）\\n王哲（东京）\\n李彧（香港）\\n黎永康（香港）\\n联系人：王有鑫\\n电话：010-66594127\\n邮件：wangyouxin_hq@bank-of-china.com主要经济体GDP增速变化趋势（%）\\n资料来源：IMF，中国银行研究院\\n\\n这个片段位于文档的开头部分，作为“全球经济金融展望报告”的摘要或简介。它概述了2023年全球经济增长、贸易和货币政策的主要趋势，并预测了2024年的经济前景。该片段还提到了几个关键的研究课题，包括海湾六国的经济发展与投资前景、美国房地产市场的脆弱性等热点问题。\\n\\n这个部分提供了报告的主题框架和主要关注点，帮助读者快速了解报告的核心内容。它通常位于文档的第一页或前言中，并且包含了报告作者团队的信息以及联系人详情，便于后续沟通和获取更多资料。\\n\\n在文档的结构上，这样的摘要部分起到了引言的作用，为详细的分析、预测和研究提供了背景信息和方向指引。通过阅读这个片段，可以快速定位到报告的主要议题和预期的研究深度，有助于读者根据自己的需求进行深入阅读或查找特定章节。')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contextualized_chunks[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f6f46c73-7369-448f-a89a-ed3d817cad47",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:38:53.304858Z",
     "iopub.status.busy": "2024-11-24T07:38:53.304699Z",
     "iopub.status.idle": "2024-11-24T07:39:09.789908Z",
     "shell.execute_reply": "2024-11-24T07:39:09.787507Z",
     "shell.execute_reply.started": "2024-11-24T07:38:53.304846Z"
    },
    "papermill": {
     "duration": 83.983138,
     "end_time": "2024-11-23T14:35:06.117207",
     "exception": false,
     "start_time": "2024-11-23T14:33:42.134069",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from langchain_community.vectorstores import Chroma\n",
    "\n",
    "model_path = 'BAAI/bge-large-zh-v1.5'\n",
    "embeddings = get_embeddings(model_path)\n",
    "\n",
    "persist_directory = os.path.join(expr_dir, 'chroma', 'bge')\n",
    "shutil.rmtree(persist_directory, ignore_errors=True)\n",
    "vector_db = Chroma.from_documents(\n",
    "    contextualized_chunks,\n",
    "    embedding=embeddings,\n",
    "    persist_directory=persist_directory\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "566c6f3c-5777-4aa9-bc60-a3ee23050506",
   "metadata": {
    "papermill": {
     "duration": 0.10028,
     "end_time": "2024-11-23T14:35:06.350323",
     "exception": false,
     "start_time": "2024-11-23T14:35:06.250043",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 计算检索准确率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b03e3382-39e9-4932-a265-69b811041629",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:09.793387Z",
     "iopub.status.busy": "2024-11-24T07:39:09.792288Z",
     "iopub.status.idle": "2024-11-24T07:39:09.805632Z",
     "shell.execute_reply": "2024-11-24T07:39:09.805210Z",
     "shell.execute_reply.started": "2024-11-24T07:39:09.793314Z"
    },
    "papermill": {
     "duration": 0.110628,
     "end_time": "2024-11-23T14:35:06.583241",
     "exception": false,
     "start_time": "2024-11-23T14:35:06.472613",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_df = qa_df[(qa_df['dataset'] == 'test') & (qa_df['qa_type'] == 'detailed')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "32c3ad14-b217-44aa-bdb9-909b9d559668",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:09.806289Z",
     "iopub.status.busy": "2024-11-24T07:39:09.806133Z",
     "iopub.status.idle": "2024-11-24T07:39:09.811727Z",
     "shell.execute_reply": "2024-11-24T07:39:09.811266Z",
     "shell.execute_reply.started": "2024-11-24T07:39:09.806277Z"
    },
    "papermill": {
     "duration": 0.138041,
     "end_time": "2024-11-23T14:35:06.823783",
     "exception": false,
     "start_time": "2024-11-23T14:35:06.685742",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_hit_stat_df(top_k_arr=list(range(1, 9))):\n",
    "    hit_stat_data = []\n",
    "\n",
    "    for k in tqdm(top_k_arr):\n",
    "        for idx, row in test_df.iterrows():\n",
    "            question = row['question']\n",
    "            true_uuid = row['uuid']\n",
    "            # chunks = retrieve_fn(question, k=k)\n",
    "            chunks = vector_db.similarity_search(question, k=k)\n",
    "            retrieved_uuids = [doc.metadata['uuid'] for doc in chunks]\n",
    "\n",
    "            hit_stat_data.append({\n",
    "                'question': question,\n",
    "                'top_k': k,\n",
    "                'hit': int(true_uuid in retrieved_uuids),\n",
    "                'retrieved_chunks': len(chunks)\n",
    "            })\n",
    "    hit_stat_df = pd.DataFrame(hit_stat_data)\n",
    "    return hit_stat_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a63797c7-4151-4f55-8e5d-080c34265393",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:09.812260Z",
     "iopub.status.busy": "2024-11-24T07:39:09.812140Z",
     "iopub.status.idle": "2024-11-24T07:39:27.228737Z",
     "shell.execute_reply": "2024-11-24T07:39:27.228386Z",
     "shell.execute_reply.started": "2024-11-24T07:39:09.812248Z"
    },
    "papermill": {
     "duration": 18.284887,
     "end_time": "2024-11-23T14:35:25.215134",
     "exception": false,
     "start_time": "2024-11-23T14:35:06.930247",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "115fd280a0fd430285874e36945c24d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "hit_stat_df = get_hit_stat_df()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d4890789-a44c-41de-b17f-0ff505788494",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.229333Z",
     "iopub.status.busy": "2024-11-24T07:39:27.229200Z",
     "iopub.status.idle": "2024-11-24T07:39:27.235789Z",
     "shell.execute_reply": "2024-11-24T07:39:27.235354Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.229320Z"
    },
    "papermill": {
     "duration": 0.145649,
     "end_time": "2024-11-23T14:35:25.499527",
     "exception": false,
     "start_time": "2024-11-23T14:35:25.353878",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>top_k</th>\n",
       "      <th>hit_rate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.419355</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>0.580645</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>0.688172</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>0.752688</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0.817204</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>6</td>\n",
       "      <td>0.827957</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>7</td>\n",
       "      <td>0.838710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>8</td>\n",
       "      <td>0.838710</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   top_k  hit_rate\n",
       "0      1  0.419355\n",
       "1      2  0.580645\n",
       "2      3  0.688172\n",
       "3      4  0.752688\n",
       "4      5  0.817204\n",
       "5      6  0.827957\n",
       "6      7  0.838710\n",
       "7      8  0.838710"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hit_stat_df.groupby(['top_k'])['hit'].mean().reset_index().rename(columns={'hit': 'hit_rate'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b0b086d1-6cec-4743-8df6-2ab3b1593689",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.236360Z",
     "iopub.status.busy": "2024-11-24T07:39:27.236233Z",
     "iopub.status.idle": "2024-11-24T07:39:27.610707Z",
     "shell.execute_reply": "2024-11-24T07:39:27.610288Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.236348Z"
    },
    "papermill": {
     "duration": 0.525548,
     "end_time": "2024-11-23T14:35:26.224632",
     "exception": false,
     "start_time": "2024-11-23T14:35:25.699084",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: xlabel='top_k', ylabel='hit'>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAjcAAAGxCAYAAACeKZf2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/TGe4hAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAorElEQVR4nO3df3RU9Z3/8ddkIAkREsSQBGIkBRQIQkITkwZqYWs0ix5WdrdutNSkI6anlbHofPULEUzEHwy2NoajlAgS8aulxHbxR1eMP6YGS4kNJtJCFZQqBIEk5FgJRJ3ozHz/8HTsLAEBk7mTT56Pc+45zJ3PzbxvPdYnd+5kbIFAICAAAABDRFk9AAAAQG8ibgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYZZDVA4Sb3+/XoUOHNGzYMNlsNqvHAQAApyEQCOjYsWMaPXq0oqJOfW1mwMXNoUOHlJaWZvUYAADgLBw4cEDnn3/+KdcMuLgZNmyYpC/+x4mPj7d4GgAAcDo6OzuVlpYW/O/4qQy4uPnHW1Hx8fHEDQAA/czp3FLCDcUAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIwyyOoBAADobW/f93urR+gVk5Z894zW33XXXX0zSJh93fPgyg0AADAKcQMAAIxC3AAAAKNwzw0AGOy+H3zP6hF6xZInf2v1COhHuHIDAACMQtwAAACjEDcAAMAoxA0AADAKNxQDGBAe/j+/s3qEXuH8xRyrRwAiHlduAACAUSyPm1WrVik9PV2xsbHKy8tTY2PjKddXVVVpwoQJGjJkiNLS0nTrrbfq008/DdO0AAAg0lkaN7W1tXK5XKqoqFBzc7MyMzNVWFio9vb2Htdv2LBBixcvVkVFhd5++22tW7dOtbW1uuOOO8I8OQAAiFSWxk1lZaVKS0vlcDiUkZGh6upqxcXFqaampsf127Zt04wZM/T9739f6enpuuKKK3Tdddd95dUeAAAwcFgWN93d3WpqalJBQcGXw0RFqaCgQA0NDT0eM336dDU1NQVj5r333tPmzZt15ZVXhmVmAAAQ+Sz7tFRHR4d8Pp+Sk5ND9icnJ2v37t09HvP9739fHR0d+va3v61AIKDPP/9cP/7xj0/5tpTX65XX6w0+7uzs7J0TAAAAEcnyG4rPRH19vZYvX65f/vKXam5u1qZNm/T888/rnnvuOekxbrdbCQkJwS0tLS2MEwMAgHCz7MpNYmKi7Ha72traQva3tbUpJSWlx2PuvPNOXX/99brxxhslSVOmTFFXV5d+9KMfacmSJYqKOrHVysrK5HK5go87OzsJHAAADGbZlZvo6GhlZ2fL4/EE9/n9fnk8HuXn5/d4zMcff3xCwNjtdklSIBDo8ZiYmBjFx8eHbAAAwFyW/oZil8ulkpIS5eTkKDc3V1VVVerq6pLD4ZAkFRcXKzU1VW63W5I0Z84cVVZWatq0acrLy9PevXt15513as6cOcHIAQAAA5ulcVNUVKQjR46ovLxcra2tysrKUl1dXfAm45aWlpArNUuXLpXNZtPSpUt18OBBjRw5UnPmzNF9991n1SkAAIAIY/l3SzmdTjmdzh6fq6+vD3k8aNAgVVRUqKKiIgyTAQCA/qhffVoKAADgqxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKIOsHgBAeG35zkyrR+gVM1/bYvUIACIUV24AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYhbgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEaJiLhZtWqV0tPTFRsbq7y8PDU2Np507axZs2Sz2U7YrrrqqjBODAAAIpXlcVNbWyuXy6WKigo1NzcrMzNThYWFam9v73H9pk2bdPjw4eC2a9cu2e12XXPNNWGeHAAARCLL46ayslKlpaVyOBzKyMhQdXW14uLiVFNT0+P6ESNGKCUlJbi9/PLLiouLI24AAIAki+Omu7tbTU1NKigoCO6LiopSQUGBGhoaTutnrFu3Ttdee63OOeecvhoTAAD0I4OsfPGOjg75fD4lJyeH7E9OTtbu3bu/8vjGxkbt2rVL69atO+kar9crr9cbfNzZ2Xn2AwMAgIhn+dtSX8e6des0ZcoU5ebmnnSN2+1WQkJCcEtLSwvjhAAAINwsjZvExETZ7Xa1tbWF7G9ra1NKSsopj+3q6tLGjRs1f/78U64rKyvT0aNHg9uBAwe+9twAACByWRo30dHRys7OlsfjCe7z+/3yeDzKz88/5bG/+c1v5PV69YMf/OCU62JiYhQfHx+yAQAAc1l6z40kuVwulZSUKCcnR7m5uaqqqlJXV5ccDockqbi4WKmpqXK73SHHrVu3TnPnztV5551nxdgAACBCWR43RUVFOnLkiMrLy9Xa2qqsrCzV1dUFbzJuaWlRVFToBaY9e/Zo69ateumll6wYGQAARDDL40aSnE6nnE5nj8/V19efsG/ChAkKBAJ9PBVMN+OhGVaP0Cv+ePMfrR4BACJKv/60FAAAwP9G3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKJbHzapVq5Senq7Y2Fjl5eWpsbHxlOs/+ugjLViwQKNGjVJMTIwuuugibd68OUzTAgCASDfIyhevra2Vy+VSdXW18vLyVFVVpcLCQu3Zs0dJSUknrO/u7tbll1+upKQk/fa3v1Vqaqr279+v4cOHh394AAAQkSyNm8rKSpWWlsrhcEiSqqur9fzzz6umpkaLFy8+YX1NTY0+/PBDbdu2TYMHD5Ykpaenh3NkAAAQ4Sx7W6q7u1tNTU0qKCj4cpioKBUUFKihoaHHY5577jnl5+drwYIFSk5O1sUXX6zly5fL5/OFa2wAABDhLLty09HRIZ/Pp+Tk5JD9ycnJ2r17d4/HvPfee/r973+vefPmafPmzdq7d69uuukmffbZZ6qoqOjxGK/XK6/XG3zc2dnZeycBAAAijuU3FJ8Jv9+vpKQkrVmzRtnZ2SoqKtKSJUtUXV190mPcbrcSEhKCW1paWhgnBgAA4WZZ3CQmJsput6utrS1kf1tbm1JSUno8ZtSoUbroootkt9uD+yZNmqTW1lZ1d3f3eExZWZmOHj0a3A4cONB7JwEAACKOZXETHR2t7OxseTye4D6/3y+Px6P8/Pwej5kxY4b27t0rv98f3PfOO+9o1KhRio6O7vGYmJgYxcfHh2wAAMBclr4t5XK5tHbtWj3++ON6++239ZOf/ERdXV3BT08VFxerrKwsuP4nP/mJPvzwQy1cuFDvvPOOnn/+eS1fvlwLFiyw6hQAAECEsfSj4EVFRTpy5IjKy8vV2tqqrKws1dXVBW8ybmlpUVTUl/2VlpamF198UbfeequmTp2q1NRULVy4UIsWLbLqFAAAQISxNG4kyel0yul09vhcfX39Cfvy8/P1+uuv9/FUAACgv+pXn5YCAAD4KsQNAAAwiuVvS8F6LXdPsXqEXnFB+U6rRwAARACu3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjBIRcbNq1Sqlp6crNjZWeXl5amxsPOna9evXy2azhWyxsbFhnBYAAEQyy+OmtrZWLpdLFRUVam5uVmZmpgoLC9Xe3n7SY+Lj43X48OHgtn///jBODAAAIpnlcVNZWanS0lI5HA5lZGSourpacXFxqqmpOekxNptNKSkpwS05OTmMEwMAgEhmadx0d3erqalJBQUFwX1RUVEqKChQQ0PDSY87fvy4xowZo7S0NF199dX661//Go5xAQBAP2Bp3HR0dMjn851w5SU5OVmtra09HjNhwgTV1NTo2Wef1ZNPPim/36/p06frgw8+6HG91+tVZ2dnyAYAAMxl+dtSZyo/P1/FxcXKysrSzJkztWnTJo0cOVKPPPJIj+vdbrcSEhKCW1paWpgnBgAA4WRp3CQmJsput6utrS1kf1tbm1JSUk7rZwwePFjTpk3T3r17e3y+rKxMR48eDW4HDhz42nMDAIDIZWncREdHKzs7Wx6PJ7jP7/fL4/EoPz//tH6Gz+fTzp07NWrUqB6fj4mJUXx8fMgGAADMNcjqAVwul0pKSpSTk6Pc3FxVVVWpq6tLDodDklRcXKzU1FS53W5J0t13361vfetbGj9+vD766CP9/Oc/1/79+3XjjTdaeRoAACBCWB43RUVFOnLkiMrLy9Xa2qqsrCzV1dUFbzJuaWlRVNSXF5j+/ve/q7S0VK2trTr33HOVnZ2tbdu2KSMjw6pTAAAAEcTyuJEkp9Mpp9PZ43P19fUhjx988EE9+OCDYZgKAAD0R/3u01IAAACnQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMclZx893vflcfffTRCfs7Ozv13e9+9+vOBAAAcNbOKm7q6+vV3d19wv5PP/1Uf/jDH772UAAAAGdr0Jks/stf/hL881tvvaXW1tbgY5/Pp7q6OqWmpvbedAAAAGfojOImKytLNptNNputx7efhgwZooceeqjXhgMAADhTZxQ377//vgKBgMaOHavGxkaNHDky+Fx0dLSSkpJkt9t7fUgAAIDTdUZxM2bMGEmS3+/vk2EAAAC+rtOOm+eee06zZ8/W4MGD9dxzz51y7b/927997cEAAADOxmnHzdy5c9Xa2qqkpCTNnTv3pOtsNpt8Pl9vzAYAAHDGTjtu/vmtKN6WAgAAkeqM7rn5Zx6PRx6PR+3t7SGxY7PZtG7dul4ZDgAA4EydVdwsW7ZMd999t3JycjRq1CjZbLbengsAAOCsnFXcVFdXa/369br++ut7ex4AAICv5ay+fqG7u1vTp0/v7VkAAAC+trOKmxtvvFEbNmzo7VkAAAC+ttN+W8rlcgX/7Pf7tWbNGr3yyiuaOnWqBg8eHLK2srKy9yYEAAA4A6cdN2+++WbI46ysLEnSrl27Qvb355uLs2//f1aP0Cuafl5s9QgAAFjmtOPm1Vdf7cs5AAAAesVZ3XMDAAAQqSIiblatWqX09HTFxsYqLy9PjY2Np3Xcxo0bZbPZTvl1EAAAYGCxPG5qa2vlcrlUUVGh5uZmZWZmqrCwUO3t7ac8bt++fbrtttt06aWXhmlSAADQH1geN5WVlSotLZXD4VBGRoaqq6sVFxenmpqakx7j8/k0b948LVu2TGPHjg3jtAAAINJZGjfd3d1qampSQUFBcF9UVJQKCgrU0NBw0uPuvvtuJSUlaf78+eEYEwAA9CNn/cWZvaGjo0M+n0/Jyckh+5OTk7V79+4ej9m6davWrVunHTt2nNZreL1eeb3e4OPOzs6znhcAAEQ+y9+WOhPHjh3T9ddfr7Vr1yoxMfG0jnG73UpISAhuaWlpfTwlAACwkqVXbhITE2W329XW1hayv62tTSkpKSes/9vf/qZ9+/Zpzpw5wX1+v1+SNGjQIO3Zs0fjxo0LOaasrCzktyt3dnYSOAAAGMzSuImOjlZ2drY8Hk/w49x+v18ej0dOp/OE9RMnTtTOnTtD9i1dulTHjh3TypUre4yWmJgYxcTE9Mn8AAAg8lgaN9IX31lVUlKinJwc5ebmqqqqSl1dXXI4HJKk4uJipaamyu12KzY2VhdffHHI8cOHD5ekE/YDAICByfK4KSoq0pEjR1ReXq7W1lZlZWWprq4ueJNxS0uLoqL61a1BAADAQpbHjSQ5nc4e34aSpPr6+lMeu379+t4fCAAA9FtcEgEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYhbgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYhbgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABglIiIm1WrVik9PV2xsbHKy8tTY2PjSddu2rRJOTk5Gj58uM455xxlZWXpiSeeCOO0AAAgklkeN7W1tXK5XKqoqFBzc7MyMzNVWFio9vb2HtePGDFCS5YsUUNDg/7yl7/I4XDI4XDoxRdfDPPkAAAgElkeN5WVlSotLZXD4VBGRoaqq6sVFxenmpqaHtfPmjVL//7v/65JkyZp3LhxWrhwoaZOnaqtW7eGeXIAABCJLI2b7u5uNTU1qaCgILgvKipKBQUFamho+MrjA4GAPB6P9uzZo+985zt9OSoAAOgnBln54h0dHfL5fEpOTg7Zn5ycrN27d5/0uKNHjyo1NVVer1d2u12//OUvdfnll/e41uv1yuv1Bh93dnb2zvAAACAiWRo3Z2vYsGHasWOHjh8/Lo/HI5fLpbFjx2rWrFknrHW73Vq2bFn4hwQAAJawNG4SExNlt9vV1tYWsr+trU0pKSknPS4qKkrjx4+XJGVlZentt9+W2+3uMW7KysrkcrmCjzs7O5WWltY7JwAAACKOpffcREdHKzs7Wx6PJ7jP7/fL4/EoPz//tH+O3+8Peevpn8XExCg+Pj5kAwAA5rL8bSmXy6WSkhLl5OQoNzdXVVVV6urqksPhkCQVFxcrNTVVbrdb0hdvM+Xk5GjcuHHyer3avHmznnjiCa1evdrK0wAAABHC8rgpKirSkSNHVF5ertbWVmVlZamuri54k3FLS4uior68wNTV1aWbbrpJH3zwgYYMGaKJEyfqySefVFFRkVWnAAAAIojlcSNJTqdTTqezx+fq6+tDHt9777269957wzAVAADojyz/JX4AAAC9ibgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYhbgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAYhbgBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFEiIm5WrVql9PR0xcbGKi8vT42NjSddu3btWl166aU699xzde6556qgoOCU6wEAwMBiedzU1tbK5XKpoqJCzc3NyszMVGFhodrb23tcX19fr+uuu06vvvqqGhoalJaWpiuuuEIHDx4M8+QAACASWR43lZWVKi0tlcPhUEZGhqqrqxUXF6eampoe1//qV7/STTfdpKysLE2cOFGPPvqo/H6/PB5PmCcHAACRyNK46e7uVlNTkwoKCoL7oqKiVFBQoIaGhtP6GR9//LE+++wzjRgxoq/GBAAA/cggK1+8o6NDPp9PycnJIfuTk5O1e/fu0/oZixYt0ujRo0MC6Z95vV55vd7g487OzrMfGAAARDzL35b6OlasWKGNGzfq6aefVmxsbI9r3G63EhISgltaWlqYpwQAAOFkadwkJibKbrerra0tZH9bW5tSUlJOeewDDzygFStW6KWXXtLUqVNPuq6srExHjx4NbgcOHOiV2QEAQGSyNG6io6OVnZ0dcjPwP24Ozs/PP+lxP/vZz3TPPfeorq5OOTk5p3yNmJgYxcfHh2wAAMBclt5zI0kul0slJSXKyclRbm6uqqqq1NXVJYfDIUkqLi5Wamqq3G63JOn+++9XeXm5NmzYoPT0dLW2tkqShg4dqqFDh1p2HgAAIDJYHjdFRUU6cuSIysvL1draqqysLNXV1QVvMm5paVFU1JcXmFavXq3u7m5973vfC/k5FRUVuuuuu8I5OgAAiECWx40kOZ1OOZ3OHp+rr68Pebxv376+HwgAAPRb/frTUgAAAP8bcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAo1geN6tWrVJ6erpiY2OVl5enxsbGk67961//qv/8z/9Uenq6bDabqqqqwjcoAADoFyyNm9raWrlcLlVUVKi5uVmZmZkqLCxUe3t7j+s//vhjjR07VitWrFBKSkqYpwUAAP2BpXFTWVmp0tJSORwOZWRkqLq6WnFxcaqpqelx/SWXXKKf//znuvbaaxUTExPmaQEAQH9gWdx0d3erqalJBQUFXw4TFaWCggI1NDRYNRYAAOjnBln1wh0dHfL5fEpOTg7Zn5ycrN27d/fa63i9Xnm93uDjzs7OXvvZAAAg8lh+Q3Ffc7vdSkhICG5paWlWjwQAAPqQZXGTmJgou92utra2kP1tbW29erNwWVmZjh49GtwOHDjQaz8bAABEHsviJjo6WtnZ2fJ4PMF9fr9fHo9H+fn5vfY6MTExio+PD9kAAIC5LLvnRpJcLpdKSkqUk5Oj3NxcVVVVqaurSw6HQ5JUXFys1NRUud1uSV/chPzWW28F/3zw4EHt2LFDQ4cO1fjx4y07DwAAEDksjZuioiIdOXJE5eXlam1tVVZWlurq6oI3Gbe0tCgq6suLS4cOHdK0adOCjx944AE98MADmjlzpurr68M9PgAAiECWxo0kOZ1OOZ3OHp/738GSnp6uQCAQhqkAAEB/ZfynpQAAwMBC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKMQNAAAwCnEDAACMQtwAAACjEDcAAMAoxA0AADAKcQMAAIxC3AAAAKMQNwAAwCjEDQAAMApxAwAAjELcAAAAoxA3AADAKBERN6tWrVJ6erpiY2OVl5enxsbGU67/zW9+o4kTJyo2NlZTpkzR5s2bwzQpAACIdJbHTW1trVwulyoqKtTc3KzMzEwVFhaqvb29x/Xbtm3Tddddp/nz5+vNN9/U3LlzNXfuXO3atSvMkwMAgEhkedxUVlaqtLRUDodDGRkZqq6uVlxcnGpqanpcv3LlSv3rv/6rbr/9dk2aNEn33HOPvvnNb+rhhx8O8+QAACASWRo33d3dampqUkFBQXBfVFSUCgoK1NDQ0OMxDQ0NIeslqbCw8KTrAQDAwDLIyhfv6OiQz+dTcnJyyP7k5GTt3r27x2NaW1t7XN/a2trjeq/XK6/XG3x89OhRSVJnZ+cJa33eT85o/kjV07mdyrFPfX00SXid6Xl//snnfTRJeJ3peXd9PjDP+xPvx300SXid6Xl/+tlnfTRJeJ3peR//tKuPJgmvMz3vf/7vXX/W03n/Y18gEPjK4y2Nm3Bwu91atmzZCfvT0tIsmCY8Eh76sdUjWMOdYPUElkhYNDDPWwkD87z/7yqrJ7DGvU8NzH/eutfqAayxYsWKkz537NgxJXzFv/+Wxk1iYqLsdrva2tpC9re1tSklJaXHY1JSUs5ofVlZmVwuV/Cx3+/Xhx9+qPPOO082m+1rnsGZ6ezsVFpamg4cOKD4+PiwvraVOG/OeyDgvDnvgcDK8w4EAjp27JhGjx79lWstjZvo6GhlZ2fL4/Fo7ty5kr6ID4/HI6fT2eMx+fn58ng8uuWWW4L7Xn75ZeXn5/e4PiYmRjExMSH7hg8f3hvjn7X4+PgB9S/DP3DeAwvnPbBw3gOLVef9VVds/sHyt6VcLpdKSkqUk5Oj3NxcVVVVqaurSw6HQ5JUXFys1NRUud1uSdLChQs1c+ZM/eIXv9BVV12ljRs36o033tCaNWusPA0AABAhLI+boqIiHTlyROXl5WptbVVWVpbq6uqCNw23tLQoKurLD3VNnz5dGzZs0NKlS3XHHXfowgsv1DPPPKOLL77YqlMAAAARxPK4kSSn03nSt6Hq6+tP2HfNNdfommuu6eOpel9MTIwqKipOeJvMdJw35z0QcN6c90DQX87bFjidz1QBAAD0E5b/hmIAAIDeRNwAAACjEDcAAMAoxE0YvPbaa5ozZ45Gjx4tm82mZ555xuqRwsLtduuSSy7RsGHDlJSUpLlz52rPnj1Wj9XnVq9eralTpwZ/D0R+fr5eeOEFq8cKuxUrVshms4X8TioT3XXXXbLZbCHbxIkTrR4rLA4ePKgf/OAHOu+88zRkyBBNmTJFb7zxhtVj9an09PQT/nnbbDYtWLDA6tH6lM/n05133qlvfOMbGjJkiMaNG6d77rnntL4KwQoR8Wkp03V1dSkzM1M33HCD/uM//sPqccJmy5YtWrBggS655BJ9/vnnuuOOO3TFFVforbfe0jnnnGP1eH3m/PPP14oVK3ThhRcqEAjo8ccf19VXX60333xTkydPtnq8sNi+fbseeeQRTZ061epRwmLy5Ml65ZVXgo8HDTL//1r//ve/a8aMGfqXf/kXvfDCCxo5cqTeffddnXvuuVaP1qe2b98un+/L7+PbtWuXLr/88n75Cd4zcf/992v16tV6/PHHNXnyZL3xxhtyOBxKSEjQT3/6U6vHO4H5/wZGgNmzZ2v27NlWjxF2dXV1IY/Xr1+vpKQkNTU16Tvf+Y5FU/W9OXPmhDy+7777tHr1ar3++usDIm6OHz+uefPmae3atbr33oHxxTiDBg066VfAmOr+++9XWlqaHnvsseC+b3zjGxZOFB4jR44MebxixQqNGzdOM2fOtGii8Ni2bZuuvvpqXXXVVZK+uIL161//Wo2NjRZP1jPelkLY/OMb2UeMGGHxJOHj8/m0ceNGdXV1nfQrQkyzYMECXXXVVSooKLB6lLB59913NXr0aI0dO1bz5s1TS0uL1SP1ueeee045OTm65pprlJSUpGnTpmnt2rVWjxVW3d3devLJJ3XDDTeE/bsKw2369OnyeDx65513JEl//vOftXXr1oj9iztXbhAWfr9ft9xyi2bMmDEgfpv0zp07lZ+fr08//VRDhw7V008/rYyMDKvH6nMbN25Uc3Oztm/fbvUoYZOXl6f169drwoQJOnz4sJYtW6ZLL71Uu3bt0rBhw6wer8+89957Wr16tVwul+644w5t375dP/3pTxUdHa2SkhKrxwuLZ555Rh999JF++MMfWj1Kn1u8eLE6Ozs1ceJE2e12+Xw+3XfffZo3b57Vo/WIuEFYLFiwQLt27dLWrVutHiUsJkyYoB07dujo0aP67W9/q5KSEm3ZssXowDlw4IAWLlyol19+WbGxsVaPEzb//DfXqVOnKi8vT2PGjNFTTz2l+fPnWzhZ3/L7/crJydHy5cslSdOmTdOuXbtUXV09YOJm3bp1mj179ml9S3V/99RTT+lXv/qVNmzYoMmTJ2vHjh265ZZbNHr06Ij8503coM85nU79z//8j1577TWdf/75Vo8TFtHR0Ro/frwkKTs7W9u3b9fKlSv1yCOPWDxZ32lqalJ7e7u++c1vBvf5fD699tprevjhh+X1emW32y2cMDyGDx+uiy66SHv37rV6lD41atSoE2J90qRJ+u///m+LJgqv/fv365VXXtGmTZusHiUsbr/9di1evFjXXnutJGnKlCnav3+/3G43cYOBJRAI6Oabb9bTTz+t+vr6AXGz4cn4/X55vV6rx+hTl112mXbu3Bmyz+FwaOLEiVq0aNGACBvpixuq//a3v+n666+3epQ+NWPGjBN+tcM777yjMWPGWDRReD322GNKSkoK3mBruo8//jjkS6wlyW63y+/3WzTRqRE3YXD8+PGQv8W9//772rFjh0aMGKELLrjAwsn61oIFC7RhwwY9++yzGjZsmFpbWyVJCQkJGjJkiMXT9Z2ysjLNnj1bF1xwgY4dO6YNGzaovr5eL774otWj9alhw4adcD/VOeeco/POO8/o+6xuu+02zZkzR2PGjNGhQ4dUUVEhu92u6667zurR+tStt96q6dOna/ny5fqv//ovNTY2as2aNVqzZo3Vo/U5v9+vxx57TCUlJQPiY//SF58Cve+++3TBBRdo8uTJevPNN1VZWakbbrjB6tF6FkCfe/XVVwOSTthKSkqsHq1P9XTOkgKPPfaY1aP1qRtuuCEwZsyYQHR0dGDkyJGByy67LPDSSy9ZPZYlZs6cGVi4cKHVY/SpoqKiwKhRowLR0dGB1NTUQFFRUWDv3r1WjxUWv/vd7wIXX3xxICYmJjBx4sTAmjVrrB4pLF588cWApMCePXusHiVsOjs7AwsXLgxccMEFgdjY2MDYsWMDS5YsCXi9XqtH6xHfCg4AAIzC77kBAABGIW4AAIBRiBsAAGAU4gYAABiFuAEAAEYhbgAAgFGIGwAAYBTiBgAAGIW4ATCgpaenq6qqyuoxAPQi4gZAxJg1a5ZuueUWq8cA0M8RNwAAwCjEDYCI8MMf/lBbtmzRypUrZbPZZLPZtG/fPm3ZskW5ubmKiYnRqFGjtHjxYn3++efB42bNmiWn0ymn06mEhAQlJibqzjvv1Nl+bd6jjz6q4cOHy+Px9NapAQgz4gZARFi5cqXy8/NVWlqqw4cP6/Dhwxo8eLCuvPJKXXLJJfrzn/+s1atXa926dbr33ntDjn388cc1aNAgNTY2auXKlaqsrNSjjz56xjP87Gc/0+LFi/XSSy/psssu661TAxBmg6weAAAkKSEhQdHR0YqLi1NKSookacmSJUpLS9PDDz8sm82miRMn6tChQ1q0aJHKy8sVFfXF38/S0tL04IMPymazacKECdq5c6cefPBBlZaWnvbrL1q0SE888YS2bNmiyZMn98k5AggPrtwAiFhvv/228vPzZbPZgvtmzJih48eP64MPPgju+9a3vhWyJj8/X++++658Pt9pvc4vfvELrV27Vlu3biVsAAMQNwAGvEsvvVQ+n09PPfWU1aMA6AXEDYCIER0dHXK1ZdKkSWpoaAi5OfiPf/yjhg0bpvPPPz+4709/+lPIz3n99dd14YUXym63n9br5ubm6oUXXtDy5cv1wAMPfM2zAGA14gZAxEhPT9ef/vQn7du3Tx0dHbrpppt04MAB3Xzzzdq9e7eeffZZVVRUyOVyBe+3kaSWlha5XC7t2bNHv/71r/XQQw9p4cKFZ/Ta06dP1+bNm7Vs2TJ+qR/Qz3FDMYCIcdttt6mkpEQZGRn65JNP9P7772vz5s26/fbblZmZqREjRmj+/PlaunRpyHHFxcX65JNPlJubK7vdroULF+pHP/rRGb/+t7/9bT3//PO68sorZbfbdfPNN/fWqQEII1vgbH8ZBABEgFmzZikrK4urLQCCeFsKAAAYhbgBYKw//OEPGjp06Ek3AGbibSkAxvrkk0908ODBkz4/fvz4ME4DIFyIGwAAYBTelgIAAEYhbgAAgFGIGwAAYBTiBgAAGIW4AQAARiFuAACAUYgbAABgFOIGAAAY5f8DMfTqsU8KN/8AAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "sns.barplot(x='top_k', y='hit', data=hit_stat_df, errorbar=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7925564a-7d30-4914-baaf-4a00abb7686d",
   "metadata": {
    "papermill": {
     "duration": 0.109216,
     "end_time": "2024-11-23T14:35:26.464009",
     "exception": false,
     "start_time": "2024-11-23T14:35:26.354793",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# 生成答案"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "50404beb-3be0-4aaa-b124-8c7a52b84531",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.613022Z",
     "iopub.status.busy": "2024-11-24T07:39:27.612710Z",
     "iopub.status.idle": "2024-11-24T07:39:27.616735Z",
     "shell.execute_reply": "2024-11-24T07:39:27.616277Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.613009Z"
    },
    "papermill": {
     "duration": 0.159318,
     "end_time": "2024-11-23T14:35:26.768506",
     "exception": false,
     "start_time": "2024-11-23T14:35:26.609188",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def rag(llm, query, n_chunks=4):\n",
    "    prompt_tmpl = \"\"\"\n",
    "你是一个金融分析师，擅长根据所获取的信息片段，对问题进行分析和推理。\n",
    "你的任务是根据所获取的信息片段（<<<<context>>><<<</context>>>之间的内容）回答问题。\n",
    "回答保持简洁，不必重复问题，不要添加描述性解释和与答案无关的任何内容。\n",
    "已知信息：\n",
    "<<<<context>>>\n",
    "{{knowledge}}\n",
    "<<<</context>>>\n",
    "\n",
    "问题：{{query}}\n",
    "请回答：\n",
    "\"\"\".strip()\n",
    "    chunks = vector_db.similarity_search(query, k=n_chunks)\n",
    "    prompt = prompt_tmpl.replace('{{knowledge}}', '\\n\\n'.join([doc.page_content for doc in chunks])).replace('{{query}}', query)\n",
    "    retry_count = 3\n",
    "\n",
    "    resp = ''\n",
    "    while retry_count > 0:\n",
    "        try:\n",
    "            resp = llm.invoke(prompt)\n",
    "            break\n",
    "        except Exception as e:\n",
    "            retry_count -= 1\n",
    "            sleeping_seconds = 2 ** (4 - retry_count)\n",
    "            print(f\"query={query}, error={e}, sleeping={sleeping_seconds}, remaining retry count={retry_count}\")\n",
    "            \n",
    "            time.sleep(sleeping_seconds)\n",
    "    \n",
    "    return resp, chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e5a804-2dc6-411c-ba71-6ccf765b2b73",
   "metadata": {
    "papermill": {
     "duration": 0.135973,
     "end_time": "2024-11-23T14:35:27.001401",
     "exception": false,
     "start_time": "2024-11-23T14:35:26.865428",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "27132c3b-0051-4df6-bf57-fd804acb8d17",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.617301Z",
     "iopub.status.busy": "2024-11-24T07:39:27.617177Z",
     "iopub.status.idle": "2024-11-24T07:39:27.684873Z",
     "shell.execute_reply": "2024-11-24T07:39:27.684424Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.617289Z"
    },
    "papermill": {
     "duration": 0.199165,
     "end_time": "2024-11-23T14:35:27.323500",
     "exception": false,
     "start_time": "2024-11-23T14:35:27.124335",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3007024/432419542.py:3: LangChainDeprecationWarning: The class `Ollama` was deprecated in LangChain 0.3.1 and will be removed in 1.0.0. An updated version of the class exists in the :class:`~langchain-ollama package and should be used instead. To use it run `pip install -U :class:`~langchain-ollama` and import as `from :class:`~langchain_ollama import OllamaLLM``.\n",
      "  ollama_llm = Ollama(\n"
     ]
    }
   ],
   "source": [
    "from langchain.llms import Ollama\n",
    "\n",
    "ollama_llm = Ollama(\n",
    "    model='qwen2:7b-instruct-32k',\n",
    "    # model='qwen2:7b-instruct',\n",
    "    base_url='http://localhost:11434',\n",
    "    top_k=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "166392d8-f801-4372-b8ad-3e79aef0b350",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.685556Z",
     "iopub.status.busy": "2024-11-24T07:39:27.685422Z",
     "iopub.status.idle": "2024-11-24T07:39:27.690451Z",
     "shell.execute_reply": "2024-11-24T07:39:27.690067Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.685544Z"
    },
    "papermill": {
     "duration": 0.141864,
     "end_time": "2024-11-23T14:35:27.564409",
     "exception": false,
     "start_time": "2024-11-23T14:35:27.422545",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "prediction_df = qa_df[qa_df['dataset'] == 'test'][['uuid', 'question', 'qa_type', 'answer']].rename(columns={'answer': 'ref_answer'})\n",
    "\n",
    "def predict(llm, prediction_df, n_chunks):\n",
    "    prediction_df = prediction_df.copy()\n",
    "    answer_dict = {}\n",
    "\n",
    "    for idx, row in tqdm(prediction_df.iterrows(), total=len(prediction_df)):\n",
    "        uuid = row['uuid']\n",
    "        question = row['question']\n",
    "        answer, chunks = rag(llm, question, n_chunks=n_chunks)\n",
    "        assert len(chunks) <= n_chunks\n",
    "        answer_dict[question] = {\n",
    "            'uuid': uuid,\n",
    "            'ref_answer': row['ref_answer'],\n",
    "            'gen_answer': answer,\n",
    "            'chunks': chunks\n",
    "        }\n",
    "        \n",
    "    prediction_df.loc[:, 'gen_answer'] = prediction_df['question'].apply(lambda q: answer_dict[q]['gen_answer'])\n",
    "    prediction_df.loc[:, 'chunks'] = prediction_df['question'].apply(lambda q: answer_dict[q]['chunks'])\n",
    "\n",
    "    return prediction_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3adcbd83-b45e-4f22-9a7f-afeb60cbf916",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.691230Z",
     "iopub.status.busy": "2024-11-24T07:39:27.690948Z",
     "iopub.status.idle": "2024-11-24T07:39:27.696513Z",
     "shell.execute_reply": "2024-11-24T07:39:27.696147Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.691217Z"
    },
    "papermill": {
     "duration": 0.115633,
     "end_time": "2024-11-23T14:35:27.785208",
     "exception": false,
     "start_time": "2024-11-23T14:35:27.669575",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_path = os.path.join(expr_dir, 'predictions.pkl')\n",
    "\n",
    "if os.path.exists(save_path):\n",
    "    pred_dfs = pickle.load(open(save_path, 'rb'))\n",
    "else:\n",
    "    pred_dfs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "ca46d5f1-e698-457d-abb6-92d83cd59c66",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:39:27.697037Z",
     "iopub.status.busy": "2024-11-24T07:39:27.696917Z",
     "iopub.status.idle": "2024-11-24T07:44:39.769063Z",
     "shell.execute_reply": "2024-11-24T07:44:39.768266Z",
     "shell.execute_reply.started": "2024-11-24T07:39:27.697025Z"
    },
    "papermill": {
     "duration": 514.92352,
     "end_time": "2024-11-23T14:44:02.805529",
     "exception": false,
     "start_time": "2024-11-23T14:35:27.882009",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b920bf2892e49c88cd7278effcbc49e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pred_df = predict(ollama_llm, prediction_df, n_chunks=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4a1c867e-d59c-454d-93f5-2c8e23b666de",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:44:39.769941Z",
     "iopub.status.busy": "2024-11-24T07:44:39.769714Z",
     "iopub.status.idle": "2024-11-24T07:44:39.774052Z",
     "shell.execute_reply": "2024-11-24T07:44:39.773255Z",
     "shell.execute_reply.started": "2024-11-24T07:44:39.769920Z"
    },
    "papermill": {
     "duration": 0.325117,
     "end_time": "2024-11-23T14:44:03.281713",
     "exception": false,
     "start_time": "2024-11-23T14:44:02.956596",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_path = os.path.join(expr_dir, 'pred_df.pkl')\n",
    "\n",
    "if not os.path.exists(save_path):\n",
    "    print(f'saving to {save_path}')\n",
    "    pickle.dump(pred_df, open(save_path, 'wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d79e974-089f-4c08-ba5e-804f6542e06a",
   "metadata": {
    "papermill": {
     "duration": 0.14423,
     "end_time": "2024-11-23T14:44:03.513124",
     "exception": false,
     "start_time": "2024-11-23T14:44:03.368894",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# 评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "217568fe-c0e4-49eb-9a7c-9fdfbc033d8a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:44:39.775232Z",
     "iopub.status.busy": "2024-11-24T07:44:39.774977Z",
     "iopub.status.idle": "2024-11-24T07:44:40.013800Z",
     "shell.execute_reply": "2024-11-24T07:44:40.013314Z",
     "shell.execute_reply.started": "2024-11-24T07:44:39.775209Z"
    },
    "papermill": {
     "duration": 0.369729,
     "end_time": "2024-11-23T14:44:04.017198",
     "exception": false,
     "start_time": "2024-11-23T14:44:03.647469",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "import time\n",
    "\n",
    "judge_llm = ChatOpenAI(\n",
    "    api_key=os.environ['LLM_API_KEY'],\n",
    "    base_url=os.environ['LLM_BASE_URL'],\n",
    "    model_name='qwen2-72b-instruct',\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "def evaluate(prediction_df):\n",
    "    \"\"\"\n",
    "    对预测结果进行打分\n",
    "    :param prediction_df: 预测结果，需要包含问题，参考答案，生成的答案，列名分别为question, ref_answer, gen_answer\n",
    "    :return 打分模型原始返回结果\n",
    "    \"\"\"\n",
    "    prompt_tmpl = \"\"\"\n",
    "你是一个经济学博士，现在我有一系列问题，有一个助手已经对这些问题进行了回答，你需要参照参考答案，评价这个助手的回答是否正确，仅回复“是”或“否”即可，不要带其他描述性内容或无关信息。\n",
    "问题：\n",
    "<question>\n",
    "{{question}}\n",
    "</question>\n",
    "\n",
    "参考答案：\n",
    "<ref_answer>\n",
    "{{ref_answer}}\n",
    "</ref_answer>\n",
    "\n",
    "助手回答：\n",
    "<gen_answer>\n",
    "{{gen_answer}}\n",
    "</gen_answer>\n",
    "请评价：\n",
    "    \"\"\"\n",
    "    results = []\n",
    "\n",
    "    for _, row in prediction_df.iterrows():\n",
    "        question = row['question']\n",
    "        ref_answer = row['ref_answer']\n",
    "        gen_answer = row['gen_answer']\n",
    "\n",
    "        prompt = prompt_tmpl.replace('{{question}}', question).replace('{{ref_answer}}', str(ref_answer)).replace('{{gen_answer}}', gen_answer).strip()\n",
    "        \n",
    "        retry_count = 3\n",
    "        result = ''\n",
    "        \n",
    "        while retry_count > 0:\n",
    "            try:\n",
    "                result = judge_llm.invoke(prompt).content\n",
    "                break\n",
    "            except Exception as e:\n",
    "                retry_count -= 1\n",
    "                sleeping_seconds = 2 ** (4 - retry_count)\n",
    "                print(f\"query={question}, error={e}, sleeping={sleeping_seconds}, remaining retry count={retry_count}\")\n",
    "                \n",
    "                time.sleep(sleeping_seconds)\n",
    "        \n",
    "        results.append(result)\n",
    "\n",
    "        time.sleep(1)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "71db81af-b8f9-47ba-958b-761896516605",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:44:40.014452Z",
     "iopub.status.busy": "2024-11-24T07:44:40.014315Z",
     "iopub.status.idle": "2024-11-24T07:47:16.009857Z",
     "shell.execute_reply": "2024-11-24T07:47:16.009298Z",
     "shell.execute_reply.started": "2024-11-24T07:44:40.014439Z"
    },
    "papermill": {
     "duration": 150.566109,
     "end_time": "2024-11-23T14:46:34.714324",
     "exception": false,
     "start_time": "2024-11-23T14:44:04.148215",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pred_df['raw_score'] = evaluate(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7da1b98e-99aa-4e11-9297-91eac1c62493",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:47:16.010512Z",
     "iopub.status.busy": "2024-11-24T07:47:16.010363Z",
     "iopub.status.idle": "2024-11-24T07:47:16.013791Z",
     "shell.execute_reply": "2024-11-24T07:47:16.013439Z",
     "shell.execute_reply.started": "2024-11-24T07:47:16.010498Z"
    },
    "papermill": {
     "duration": 0.138037,
     "end_time": "2024-11-23T14:46:35.040595",
     "exception": false,
     "start_time": "2024-11-23T14:46:34.902558",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['是', '否'], dtype=object)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_df['raw_score'].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2c99c078-d294-40b8-b57b-31cfd7349c3e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:47:16.014527Z",
     "iopub.status.busy": "2024-11-24T07:47:16.014260Z",
     "iopub.status.idle": "2024-11-24T07:47:16.026970Z",
     "shell.execute_reply": "2024-11-24T07:47:16.026536Z",
     "shell.execute_reply.started": "2024-11-24T07:47:16.014514Z"
    },
    "papermill": {
     "duration": 0.107466,
     "end_time": "2024-11-23T14:46:35.243603",
     "exception": false,
     "start_time": "2024-11-23T14:46:35.136137",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "pred_df['score'] = (pred_df['raw_score'] == '是').astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "423897f2-786e-415b-a613-55a4359faf76",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:47:16.027630Z",
     "iopub.status.busy": "2024-11-24T07:47:16.027480Z",
     "iopub.status.idle": "2024-11-24T07:47:16.033060Z",
     "shell.execute_reply": "2024-11-24T07:47:16.032718Z",
     "shell.execute_reply.started": "2024-11-24T07:47:16.027616Z"
    },
    "papermill": {
     "duration": 0.094328,
     "end_time": "2024-11-23T14:46:35.431162",
     "exception": false,
     "start_time": "2024-11-23T14:46:35.336834",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.68"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_df['score'].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "79325429-9cf1-4e2c-95ac-cb0c1a3b6156",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T07:47:16.033645Z",
     "iopub.status.busy": "2024-11-24T07:47:16.033509Z",
     "iopub.status.idle": "2024-11-24T07:47:16.039893Z",
     "shell.execute_reply": "2024-11-24T07:47:16.039393Z",
     "shell.execute_reply.started": "2024-11-24T07:47:16.033632Z"
    },
    "papermill": {
     "duration": 0.289336,
     "end_time": "2024-11-23T14:46:35.804651",
     "exception": false,
     "start_time": "2024-11-23T14:46:35.515315",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_path = os.path.join(expr_dir, 'eval_df.pkl')\n",
    "\n",
    "if not os.path.exists(save_path):\n",
    "    print(f'saving to {save_path}')\n",
    "    pickle.dump(pred_df, open(save_path, 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88fc7227-9c21-48da-b179-5070406eb113",
   "metadata": {
    "papermill": {
     "duration": 0.088622,
     "end_time": "2024-11-23T14:46:36.016801",
     "exception": false,
     "start_time": "2024-11-23T14:46:35.928179",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 1058.563616,
   "end_time": "2024-11-23T14:46:37.625874",
   "environment_variables": {},
   "exception": null,
   "input_path": "13_contextual_embeddings.ipynb",
   "output_path": "run_13_contextual_embeddings.ipynb",
   "parameters": {},
   "start_time": "2024-11-23T14:28:59.062258",
   "version": "2.6.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "0cd8c168767249f2a5fa412173f6e751": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_5ce1d1d9d86c40d9839877ff95734491",
       "max": 100,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_231702cf4d79477f9d5548665a1b18fe",
       "tabbable": null,
       "tooltip": null,
       "value": 100
      }
     },
     "2133bb8d85d34b8db112b4408ad60320": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "231702cf4d79477f9d5548665a1b18fe": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "23b1ad9c0f9c46c888da66e85c90eb84": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "24e6eadc3dc940ecabf30dd1a3c6d1f3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_fa4bddf2c33241b5bf918054518f128f",
       "max": 52,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_edc33e82be8f41eba6a18a0ef074ab7a",
       "tabbable": null,
       "tooltip": null,
       "value": 52
      }
     },
     "2f60367b1c8941e2bf71661c33969ae8": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3865f25c78aa46f29a25d807205281c3": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3d0b06deaa654b989eece8cde06fa0f8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "3f8ceda83287475b97608e42f5f6782f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "4881e496f1c84fe29ce9ebebaddfb3c2": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_bd096d5d219a467786a85cfe1613fedd",
        "IPY_MODEL_24e6eadc3dc940ecabf30dd1a3c6d1f3",
        "IPY_MODEL_bc2b8104b4244d8cacedeb95e800d91c"
       ],
       "layout": "IPY_MODEL_6b9a8e43c1c342dba500a14e7149b600",
       "tabbable": null,
       "tooltip": null
      }
     },
     "5ce1d1d9d86c40d9839877ff95734491": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "5ddb08be5cc64c9ab40a1d62a21763a5": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_86283159049d48b1adcfb2de2d404d4d",
       "placeholder": "​",
       "style": "IPY_MODEL_2133bb8d85d34b8db112b4408ad60320",
       "tabbable": null,
       "tooltip": null,
       "value": " 100/100 [08:34&lt;00:00, 10.01s/it]"
      }
     },
     "5ef9d83ccad1471f85335900a24a8553": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "6b9a8e43c1c342dba500a14e7149b600": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "816a079a8c804fbfa9b9a74f941abea8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_bcc69ec5db1b4aab977807284c9290e7",
        "IPY_MODEL_0cd8c168767249f2a5fa412173f6e751",
        "IPY_MODEL_5ddb08be5cc64c9ab40a1d62a21763a5"
       ],
       "layout": "IPY_MODEL_d1178c6858284f788a80b5f2a14fd0b7",
       "tabbable": null,
       "tooltip": null
      }
     },
     "86283159049d48b1adcfb2de2d404d4d": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "8ff8262c56604119883f4a5f13bb74ab": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_5ef9d83ccad1471f85335900a24a8553",
       "placeholder": "​",
       "style": "IPY_MODEL_e89e77133c344fc48c1d62f5a607ec93",
       "tabbable": null,
       "tooltip": null,
       "value": " 8/8 [00:18&lt;00:00,  2.27s/it]"
      }
     },
     "9189a076554543aaa6f5ee04e40dbe1b": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "988e6697a2af486fadeaf0b84347b565": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_e1aae4c55cb64f379e74f15357275628",
        "IPY_MODEL_fd9e23198ca1489a9773fda3510bf857",
        "IPY_MODEL_8ff8262c56604119883f4a5f13bb74ab"
       ],
       "layout": "IPY_MODEL_d2ee15001d2244529f7e47d3333c0f8e",
       "tabbable": null,
       "tooltip": null
      }
     },
     "9fc7d91f94a94933bde5ba80e64587de": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "a7d240a289084bdfba4724c0efd5ab07": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "bc2b8104b4244d8cacedeb95e800d91c": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_2f60367b1c8941e2bf71661c33969ae8",
       "placeholder": "​",
       "style": "IPY_MODEL_9fc7d91f94a94933bde5ba80e64587de",
       "tabbable": null,
       "tooltip": null,
       "value": " 52/52 [04:26&lt;00:00,  4.22s/it]"
      }
     },
     "bcc69ec5db1b4aab977807284c9290e7": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3f8ceda83287475b97608e42f5f6782f",
       "placeholder": "​",
       "style": "IPY_MODEL_3d0b06deaa654b989eece8cde06fa0f8",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "bd096d5d219a467786a85cfe1613fedd": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3865f25c78aa46f29a25d807205281c3",
       "placeholder": "​",
       "style": "IPY_MODEL_9189a076554543aaa6f5ee04e40dbe1b",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "cc3ed8dc4a5c43aca7b62d904865b2fa": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "cf68b6fe24964ce792aa63827489cb97": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "d1178c6858284f788a80b5f2a14fd0b7": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "d2ee15001d2244529f7e47d3333c0f8e": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e1aae4c55cb64f379e74f15357275628": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_23b1ad9c0f9c46c888da66e85c90eb84",
       "placeholder": "​",
       "style": "IPY_MODEL_cf68b6fe24964ce792aa63827489cb97",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "e89e77133c344fc48c1d62f5a607ec93": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "edc33e82be8f41eba6a18a0ef074ab7a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "fa4bddf2c33241b5bf918054518f128f": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "fd9e23198ca1489a9773fda3510bf857": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_cc3ed8dc4a5c43aca7b62d904865b2fa",
       "max": 8,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_a7d240a289084bdfba4724c0efd5ab07",
       "tabbable": null,
       "tooltip": null,
       "value": 8
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
